from vit import ViT

from PIL import Image
import torch
from utils.vit_utils import get_transform
import numpy as np

if __name__ == '__main__':
    mode = "predict"
    image_path = "model_data/test/plane1.jpg"
    vit = ViT()
    if mode == "predict":
        transform = get_transform(mode)
        image = Image.open(image_path)
        detect_image = vit.detect_image(image, transform)
        # TODO 是否需要后续处理
        print(detect_image)
